from typing import List, Dict, Any
from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import (
    BaseModel,
    Field,
    root_validator
)

from .llm_api import RestAPI
from db.models.llm_profile.crud import get_llm_model_by_id

requests_url = get_llm_model_by_id(llm_model_id=1)["requests_url"]


class TaiChiEmbeddings(BaseModel, Embeddings):
    """定义向量模型"""

    client: Any = None  # 访问AI的客户端

    base_url: str = None  # 访问AI的服务器地址

    api_key: str = None  # 访问AI的API密钥

    model: str = Field(default="embedding-2")  # 访问AI的embedding模型

    @root_validator
    def validate_environment(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        base_url = values["base_url"]
        api_key = values["api_key"]
        values["client"] = RestAPI(base_url=base_url, api_key=api_key)
        """验证环境"""
        return values

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed search docs"""
        return [self._get_embedding(text) for text in texts]

    def embed_query(self, text: str) -> List[float]:
        """Embed query text."""
        return self._get_embedding(text)

    def _get_embedding(self, text: str) -> List[float]:
        response = self.client.action_post(
            request=requests_url,
            model=self.model,
            input=text,
        )
        if "data" in response and response["data"]:
            results = response["data"]
            if len(results) == 1 and "embedding" in results[0]:
                return results[0]["embedding"]
            else:
                raise ValueError("结果数量不符合预期或结果中缺少 'embedding'")
        else:
            raise Exception(f"API 响应错误: {response.get('error', '没有错误消息')}")
